package com.explorati.LeetCode347.topKFrequent;

import java.util.Comparator;
import java.util.HashMap;
import java.util.Map;
import java.util.PriorityQueue;

public class Solution1 {

    /**
     * 给你一个整数数组 nums 和一个整数 k ，请你返回其中出现频率前 k 高的元素。你可以按 任意顺序 返回答案。
     *
     * 输入: nums = [1,1,1,2,2,3], k = 2
     * 输出: [1,2]
     *
     * 输入: nums = [1], k = 1
     * 输出: [1]
     *
     * @param nums
     * @param k
     * @return
     */
    public int[] topKFrequent(int[] nums, int k) {
        Map<Integer, Integer> map = new HashMap<>();
        for(int num : nums) {
            map.put(num, map.getOrDefault(num, 0) + 1);
        }

        //最小堆 int[] 的第一个元素代表数组的值，第二个元素代表了该值出现的次数
        PriorityQueue<int[]> pq = new PriorityQueue<int[]>(new Comparator<int[]>() {
            @Override
            public int compare(int[] o1, int[] o2) {
                return o1[1] - o2[1];
            }
        });

        for(Map.Entry entry : map.entrySet()) {
            int num = (int) entry.getKey(); int freq = (int) entry.getValue();

            if(pq.size() == k) {
                if(pq.peek()[1] < freq) {
                    pq.poll();
                    pq.offer(new int[]{num, freq});
                }
            }else {
                pq.offer(new int[]{num, freq});
            }
        }
        int[] ret = new int[k];
        for (int i = 0; i < k; ++i) {
            ret[i] = pq.poll()[0];
        }

        return ret;
    }

    public static void main(String[] args) {
        int[] nums = new int[]{5,2,5,3,5,3,1,1,3};
        int k = 2;
        Solution1 solution = new Solution1();

        int[] solutions = solution.topKFrequent(nums, k);
        for(int num : solutions) {
            System.out.println(num);
        }
    }
}
